# 根据xml文件裁剪图像中的特定类别，并保存到指定位置

from __future__ import division
import os
from PIL import Image
import xml.dom.minidom
import numpy as np
ImgPath = '/media/liuyi/024b9406-a087-42e6-8444-51e6a224c28e/Downloads/obj_det_vehicle/vehicle/VOC2007/JPEGImages/'
AnnoPath = '/media/liuyi/024b9406-a087-42e6-8444-51e6a224c28e/Downloads/obj_det_vehicle/vehicle/VOC2007/Annotations/'
save_Path = '/media/liuyi/024b9406-a087-42e6-8444-51e6a224c28e/Downloads/cut/'
imagelist = os.listdir(ImgPath)
for image in imagelist:
    image_pre, ext = os.path.splitext(image)
    imgfile = ImgPath + image
    print(imgfile)
    if not os.path.exists(AnnoPath + image_pre + '.xml'): continue
    xmlfile = AnnoPath + image_pre + '.xml'
    DomTree = xml.dom.minidom.parse(xmlfile)
    annotation = DomTree.documentElement
    filenamelist = annotation.getElementsByTagName('filename')  # [<DOM Element: filename at 0x381f788>]
    # filename = filenamelist[0].childNodes[0].data
    objectlist = annotation.getElementsByTagName('object')
    i = 1
    for objects in objectlist:
        namelist = objects.getElementsByTagName('name')
        objectname = namelist[0].childNodes[0].data
        if objectname == 'motorback':  # 选择特定类别
            savepath = save_Path + objectname
            if not os.path.exists(savepath):
              os.makedirs(savepath)
            bndbox = objects.getElementsByTagName('bndbox')
            cropboxes = []
            for box in bndbox:
                x1_list = box.getElementsByTagName('xmin')
                x1 = int(x1_list[0].childNodes[0].data)
                y1_list = box.getElementsByTagName('ymin')
                y1 = int(y1_list[0].childNodes[0].data)
                x2_list = box.getElementsByTagName('xmax')
                x2 = int(x2_list[0].childNodes[0].data)
                y2_list = box.getElementsByTagName('ymax')
                y2 = int(y2_list[0].childNodes[0].data)
                w = x2 - x1
                h = y2 - y1
                obj = np.array([x1, y1, x2, y2])
                shift = np.array([[1, 1, 1, 1]])
                XYmatrix = np.tile(obj, (1, 1))
                cropboxes = XYmatrix * shift
                img = Image.open(imgfile)
                if w > img.size[0] or h > img.size[1]: continue
                for cropbox in cropboxes:
                    cropedimg = img.crop(cropbox)
                    cropedimg.save(savepath + '/' + image_pre + '_' + str(i) + '.jpg')
                    i += 1

